import os
import errno
import numpy as np

from copy import deepcopy
from miscc.pt_config import cfg

from torch.nn import init
import torch
import torch.nn as nn
import torchvision.utils as vutils
import PIL.Image

#############################
def KL_loss(mu, logvar):
    """
    计算KL度量
    """
    # -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
    KLD = torch.mean(KLD_element).mul_(-0.5)
    return KLD


def compute_discriminator_loss(netD, real_imgs, fake_imgs,
                               real_labels, fake_labels,
                               conditions, gpus):
    """
    计算判别器的损失
    """
    criterion = nn.BCELoss()
    batch_size = real_imgs.size(0)
    # detach用于将tensor分离出来，就是不参与梯度计算
    cond = conditions.detach()
    fake = fake_imgs.detach()
    real_features = nn.parallel.data_parallel(netD, (real_imgs), gpus)
    fake_features = nn.parallel.data_parallel(netD, (fake), gpus)
    # real pairs
    # 当embedding和图片相符合且图片为真实图片而非生成图片时，判别器应靠近1
    inputs = (real_features, cond)
    real_logits = nn.parallel.data_parallel(netD.get_cond_logits, inputs, gpus)
    errD_real = criterion(real_logits, real_labels)
    # wrong pairs
    # 当embedding和图片不相符且图片为真实图片而非生成图片时，判别器应靠近0
    # 这里也就是让这个判别器兼顾到embedding信息
    inputs = (real_features[:(batch_size-1)], cond[1:])
    wrong_logits = \
        nn.parallel.data_parallel(netD.get_cond_logits, inputs, gpus)
    errD_wrong = criterion(wrong_logits, fake_labels[1:])
    # 当embedding和图片相符合且图片为生成图片而非真实图片时，判别器应靠近0
    # fake pairs
    inputs = (fake_features, cond)
    fake_logits = nn.parallel.data_parallel(netD.get_cond_logits, inputs, gpus)
    errD_fake = criterion(fake_logits, fake_labels)

    # 如果是阶段二的判别器，其实这也是体现阶段一判别器和阶段二判别器区别的一点，就是
    # 阶段一判别器需要确保图片来自真实描述，而阶段二判别器就不仅需要关注数据是不是真实描述的，因为阶段一判别器
    # 已经做了这个工作，阶段二判别器更需要的是判定图片的真实性
    if netD.get_uncond_logits is not None:
        real_logits = \
            nn.parallel.data_parallel(netD.get_uncond_logits,
                                      (real_features), gpus)
        fake_logits = \
            nn.parallel.data_parallel(netD.get_uncond_logits,
                                      (fake_features), gpus)
        
        uncond_errD_real = criterion(real_logits, real_labels)
        uncond_errD_fake = criterion(fake_logits, fake_labels)
        #
        errD = ((errD_real + uncond_errD_real) / 2. +
                (errD_fake + errD_wrong + uncond_errD_fake) / 3.)
        errD_real = (errD_real + uncond_errD_real) / 2.
        errD_fake = (errD_fake + uncond_errD_fake) / 2.
    else:
        errD = errD_real + (errD_fake + errD_wrong) * 0.5
    return errD, errD_real, errD_wrong, errD_fake
    # return errD, errD_real.data[0], errD_wrong.data[0], errD_fake.data[0]


def compute_generator_loss(netD, fake_imgs, real_labels, conditions, gpus):
    """
    生成器的损失
    """
    criterion = nn.BCELoss()
    cond = conditions.detach()
    fake_features = nn.parallel.data_parallel(netD, (fake_imgs), gpus)
    # fake pairs
    # 迷惑判别器，让他觉得生成图片是真实图片
    inputs = (fake_features, cond)
    fake_logits = nn.parallel.data_parallel(netD.get_cond_logits, inputs, gpus)
    errD_fake = criterion(fake_logits, real_labels)
    if netD.get_uncond_logits is not None:
        fake_logits = \
            nn.parallel.data_parallel(netD.get_uncond_logits,
                                      (fake_features), gpus)
        uncond_errD_fake = criterion(fake_logits, real_labels)
        errD_fake += uncond_errD_fake
    return errD_fake


#############################
def weights_init(m):
    """
    用于初始化模型的权值和偏置
    """
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)
    elif classname.find('Linear') != -1:
        m.weight.data.normal_(0.0, 0.02)
        if m.bias is not None:
            m.bias.data.fill_(0.0)


#############################
def save_img_results(data_img, fake, epoch, image_dir):
    """
    保存生成的图片结果，如果是阶段一的结果，那么保存原图和阶段一生成图，如果是阶段二的结果，那么只保存阶段二生成图，
    实际上，如果结合trainer代码看，如果是阶段二的结果的话，由于调用了两次此函数，实际上会保存原图，阶段一生成图和阶段二生成图
    """
    print(np.shape(fake))
    num = cfg.VIS_COUNT
    fake = fake[0:num]
    # data_img is changed to [0,1]
    if data_img is not None:
        data_img = data_img[0:num]
        vutils.save_image(
            data_img, '%s/real_samples.png' % image_dir,
            normalize=True)
        # fake.data is still [-1, 1]
        vutils.save_image(
            fake.data, '%s/fake_samples_epoch_%03d.png' %
            (image_dir, epoch), normalize=True)
    else:
        vutils.save_image(
            fake.data, '%s/lr_fake_samples_epoch_%03d.png' %
            (image_dir, epoch), normalize=True)


def save_model(netG, netD, epoch, model_dir):
    """
    保存模型
    """
    torch.save(
        netG.state_dict(),
        '%s/netG_epoch_%d.pth' % (model_dir, epoch))
    # torch.save(
    #     netD.state_dict(),
    #     '%s/netD_epoch_last.pth' % (model_dir))
    print('Save G/D models')


def mkdir_p(path):
    """
    创造文件夹
    """
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise


def save_img_predicts(fakes, stage, image_dir):
    """
    保存生成的图片结果，如果是阶段一的结果，那么保存原图和阶段一生成图，如果是阶段二的结果，那么只保存阶段二生成图，
    实际上，如果结合trainer代码看，如果是阶段二的结果的话，由于调用了两次此函数，实际上会保存原图，阶段一生成图和阶段二生成图
    """
    # print(np.shape(fakes))
    # fakes_arr = fakes.detach().cpu().numpy()
    # fakes_arr = fakes_arr.reshape((len(fakes_arr),256,256,3))
    # # fakes_arr = np.array(fakes)
    # for i in range(len(fakes_arr)):
    #     print(fakes_arr[i].shape)
        
    #     image = PIL.Image.fromarray(np.uint8(fakes_arr[i])).convert("RGB")
    #     image.save('%s/stage_%d_fake_predicted_image_%d.png' %
    #     (image_dir, stage, i))
    print(np.shape(fakes))
    # for i in range(len(fakes)):
    #     vutils.save_image(
    #         fakes[i], 
    #         '%s/stage_%d_fake_predicted_image_%d.png' % (image_dir, stage, i),
    #         normalize=True)

    for i in range(len(fakes)):
        vutils.save_image(
            fakes[i], 
            '%s/stage_%d_fake_predicted_image_%d.png' % (image_dir, stage, i),
            normalize=True)
        if image_dir == "static":
            vutils.save_image(
                fakes[i],
                '%s/%d.jpg' % (image_dir, i+2),
                normalize=True)
